from InteractScore import InteractScore
import numpy as np
import torch
import torch.utils.data as Data
from FileLoader import FileLoader
from InteractScore import InteractScore
from tqdm import tqdm


class DataLoader:
    def __init__(self, hyper_params):
        self.hyper_params = hyper_params

        # Load from file
        self.file_loader = FileLoader(hyper_params)
        self.interact_score = InteractScore(hyper_params)

        self.neg_sample_cnt = hyper_params['neg_sample_cnt']
        self.epsilon = hyper_params['epsilon']

    def generate_dataset(self, save_dataset=True) -> Data.Dataset:
        try:
            dataset = torch.load('model_dat/dataset.pkl')
            print('Restore from model_dat')
            return dataset
        except Exception as err:
            print(err)

        print('Generating Dataset...')
        positive_item = []
        negative_item = []
        item_weight = []
        user_id = []
        positive_extra = []
        negative_extra = []

        for uid, interact_hist in tqdm(self.file_loader.intereaction_list.items()):
            # Get Positive Negative Pairs
            score_dict, extra_dict = self.interact_score.get_interact_score(
                interact_hist)

            score_list = sorted(score_dict.items(),
                                key=lambda x: x[1], reverse=True)

            # Phase 1
            book_set = set(score_dict.keys())
            not_book_set = set(range(self.file_loader.book_cnt)) - book_set
            book_len = len(score_list)

            for i in range(book_len-1):
                for j in range(i+1, book_len):
                    if score_list[i][1] - score_list[j][1] > self.epsilon:
                        # Add To Ds
                        positive_item.append(score_list[i][0])
                        negative_item.append(score_list[j][0])
                        item_weight.append(score_list[i][1] - score_list[j][1])
                        user_id.append(uid)
                        positive_extra.append(extra_dict[score_list[i][0]])
                        negative_extra.append(extra_dict[score_list[j][0]])

            # Phase 2
            book_list = []
            for book in book_set:
                if score_dict[book] > self.epsilon:
                    book_list.append(book)

            not_book_list = list(not_book_set)

            phase2_positive = np.random.choice(
                book_list, size=self.neg_sample_cnt)
            phase2_negative = np.random.choice(
                not_book_list, size=self.neg_sample_cnt)
            phase2_weight = list(map(score_dict.get, phase2_positive))

            positive_item.extend(phase2_positive)
            negative_item.extend(phase2_negative)
            item_weight.extend(phase2_weight)
            user_id.extend([uid] * self.neg_sample_cnt)
            positive_extra.extend([extra_dict[i] for i in phase2_positive])
            negative_extra.extend(
                [torch.zeros(self.hyper_params['extra_dim'], dtype=torch.float32)]*self.neg_sample_cnt)

        positive_tensor = torch.LongTensor(positive_item)
        negative_tensor = torch.LongTensor(negative_item)
        weight_tensor = torch.FloatTensor(item_weight)
        user_id_tensor = torch.LongTensor(user_id)
        positive_extra_tensor = torch.stack(positive_extra, dim=0)
        negative_extra_tensor = torch.stack(negative_extra, dim=0)

        dataset = Data.TensorDataset(
            positive_tensor, negative_tensor, weight_tensor, user_id_tensor, positive_extra_tensor, negative_extra_tensor)

        if save_dataset:
            torch.save(dataset, 'model_dat/dataset.pkl')
        print('Finished!')

        return dataset
